"""
QualityJFactor Strategy Implementation for Zipline
Focuses on quality factors (ROE, gross margin, debt ratio)
"""

from zipline.api import order_target_percent, symbol, get_datetime
from zipline.finance import commission, slippage
import pandas as pd

def initialize(context):
    # Strategy parameters
    context.roe_threshold = 0.15  # 15%
    context.gross_margin_threshold = 0.3  # 30%
    context.debt_ratio_threshold = 0.6  # 60%
    context.quality_percentile = 0.2  # Top 20%
    context.max_positions = 35
    context.max_position_size = 0.04  # 4% per stock
    context.stop_loss_daily = 0.02  # 2% daily stop loss
    context.stop_loss_monthly = 0.08  # 8% monthly stop loss
    
    # Rebalancing tracking
    context.last_rebalance = None
    context.last_quarter = None
    context.last_month = None
    
    # Set commission and slippage
    context.set_commission(commission.PerShare(cost=0.005, min_trade_cost=1))
    context.set_slippage(slippage.FixedSlippage(spread=0.005))

def handle_data(context, data):
    current_date = get_datetime().normalize()
    
    # Initialize dates if not set
    if context.last_rebalance is None:
        context.last_rebalance = current_date
    if context.last_quarter is None:
        context.last_quarter = current_date
    if context.last_month is None:
        context.last_month = current_date
    
    # Semi-annual rebalancing (every 6 months)
    if (current_date - context.last_rebalance).days >= 180:
        rebalance(context, data)
        context.last_rebalance = current_date
    
    # Quarterly rebalancing to equal weights
    current_quarter = current_date.quarter
    if context.last_quarter != current_quarter:
        rebalance_to_equal_weight(context, data)
        context.last_quarter = current_quarter
    
    # Monthly return tracking
    current_month = current_date.month
    if context.last_month != current_month:
        context.portfolio.starting_value = context.portfolio.portfolio_value
        context.last_month = current_month
    
    # Daily stop loss check
    stop_loss_check(context, data)

def rebalance(context, data):
    """Semi-annual rebalancing based on quality factors"""
    # Get fundamental data and calculate quality scores
    universe = get_fundamental_data(context, data)
    quality_scores = calculate_quality_scores(context, data, universe)
    selected = filter_stocks(context, universe, quality_scores)
    
    # Exit positions not in selected stocks
    for stock in context.portfolio.positions:
        if stock not in selected:
            order_target_percent(stock, 0)
    
    # Allocate equal weights to selected stocks
    weight = min(context.max_position_size, 1.0 / len(selected))
    for stock in selected:
        order_target_percent(stock, weight)

def rebalance_to_equal_weight(context, data):
    """Quarterly rebalancing to equal weights"""
    positions = context.portfolio.positions
    if len(positions) == 0:
        return
    
    # Calculate equal weight
    weight = 1.0 / len(positions)
    
    # Rebalance existing positions
    for stock in positions:
        order_target_percent(stock, weight)

def get_fundamental_data(context, data):
    """Retrieve stocks meeting fundamental criteria"""
    # In production, replace with actual fundamental data source
    # This is a placeholder implementation
    universe = [
        stock for stock in data.fetcher_assets
        if data.current(stock, 'roe') > context.roe_threshold
        and data.current(stock, 'gross_margin') > context.gross_margin_threshold
        and data.current(stock, 'debt_ratio') < context.debt_ratio_threshold
    ]
    return universe

def calculate_quality_scores(context, data, universe):
    """Calculate quality scores for stocks"""
    scores = {}
    for stock in universe:
        # Simple quality score: average of normalized factors
        roe = data.current(stock, 'roe')
        gross_margin = data.current(stock, 'gross_margin')
        debt_ratio = 1 - data.current(stock, 'debt_ratio')  # Inverse since lower debt is better
        
        # Normalize and average
        score = (roe + gross_margin + debt_ratio) / 3
        scores[stock] = score
    return scores

def filter_stocks(context, universe, scores):
    """Filter stocks based on quality score threshold"""
    # Sort by quality score and select top 20%
    sorted_stocks = sorted(universe, key=lambda x: scores[x], reverse=True)
    top_count = max(1, int(len(sorted_stocks) * context.quality_percentile))
    return sorted_stocks[:min(top_count, context.max_positions)]

def stop_loss_check(context, data):
    """Implement daily and monthly stop loss logic"""
    portfolio = context.portfolio
    
    # Calculate daily returns
    daily_return = (portfolio.returns[-1] if len(portfolio.returns) > 0 else 0)
    
    # Calculate monthly returns
    monthly_return = (portfolio.portfolio_value / context.portfolio.starting_cash) - 1
    
    # Implement stop loss rules
    if daily_return < -context.stop_loss_daily:
        reduce_risk_exposure(context, data)
    if monthly_return < -context.stop_loss_monthly:
        liquidate_positions(context, data)

def reduce_risk_exposure(context, data):
    """Reduce position sizes by 50%"""
    for stock in context.portfolio.positions:
        current_position = context.portfolio.positions[stock].amount
        target_position = current_position * 0.5
        order_target_percent(stock, target_position)

def liquidate_positions(context, data):
    """Liquidate all positions"""
    for stock in context.portfolio.positions:
        order_target_percent(stock, 0)